import sys
import numpy as np
import pandas as pd
from PyQt5.QtWidgets import (QApplication, QMainWindow, QVBoxLayout, QWidget,
                             QHBoxLayout, QPushButton, QLabel)
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


class ThreeDPlotWidget(FigureCanvas):
    def __init__(self, parent=None, width=10, height=8, dpi=100):
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        super().__init__(self.fig)
        self.setParent(parent)

        # 创建 3D 坐标轴
        self.ax = self.fig.add_subplot(111, projection='3d')

    def plot_well_trajectory_with_dogleg(self, df, use_dogleg_severity=True):
        """绘制井眼轨迹 3D 图，可选择使用狗腿度作为颜色映射"""
        # 清除之前的图形
        self.ax.clear()

        # 提取数据
        north = df['北移'].values
        east = df['东移'].values
        depth = df['井垂深'].values

        if use_dogleg_severity:
            # 使用狗腿度作为颜色映射
            colors = df['狗腿度'].values
            cmap = 'hot'  # 使用热力图颜色，红色表示高狗腿度
            color_label = '狗腿度 (°/30m)'
            title_suffix = ' - 颜色表示狗腿度'
        else:
            # 使用井斜深作为颜色映射（原始方式）
            colors = df['井斜深'].values
            cmap = 'viridis'
            color_label = '井斜深 (m)'
            title_suffix = ''

        # 绘制 3D 轨迹线
        scatter = self.ax.scatter(east, north, -depth,
                                  c=colors, cmap=cmap,
                                  s=20, alpha=0.8)

        # 绘制连线
        self.ax.plot(east, north, -depth, 'k-', alpha=0.5, linewidth=1)

        # 标记高狗腿度区域（如果使用狗腿度）
        if use_dogleg_severity and len(df) > 0:
            high_dls_threshold = np.percentile(df['狗腿度'].values, 90)
            high_dls_mask = df['狗腿度'] > high_dls_threshold

            if np.any(high_dls_mask):
                self.ax.scatter(east[high_dls_mask], north[high_dls_mask], -depth[high_dls_mask],
                                color='red', s=50, alpha=0.9, label=f'高狗腿度区域(>{high_dls_threshold:.1f}°/30m)')

        # 标记起点和终点
        self.ax.scatter(east[0], north[0], -depth[0],
                        color='green', s=100, label='起点', marker='o')
        self.ax.scatter(east[-1], north[-1], -depth[-1],
                        color='blue', s=100, label='终点', marker='s')

        # 设置坐标轴标签
        self.ax.set_xlabel('东移 (m)')
        self.ax.set_ylabel('北移 (m)')
        self.ax.set_zlabel('垂深 (m)')

        # 设置标题
        self.ax.set_title(f'井眼轨迹 3D 图{title_suffix}')

        # 添加颜色条
        cbar = self.fig.colorbar(scatter, ax=self.ax, shrink=0.6, aspect=20)
        cbar.set_label(color_label)

        # 添加图例
        self.ax.legend()

        # 自动调整视角
        self.ax.view_init(elev=20, azim=45)

        self.draw()


class MainWindow(QMainWindow):
    def __init__(self, df):
        super().__init__()
        self.df = df
        self.use_dogleg_severity = True  # 默认使用狗腿度
        self.initUI()

    def initUI(self):
        self.setWindowTitle('井眼轨迹 3D 可视化')
        self.setGeometry(100, 100, 1200, 800)

        # 创建中央部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        # 创建布局
        layout = QVBoxLayout(central_widget)

        # 创建控制面板
        control_layout = QHBoxLayout()

        # 添加狗腿度统计信息
        if '狗腿度' in self.df.columns:
            max_dogleg = self.df['狗腿度'].max()
            avg_dogleg = self.df['狗腿度'].mean()
            high_dogleg_count = len(self.df[self.df['狗腿度'] > 5])  # 假设5°/30m为高狗腿度

            stats_label = QLabel(f"最大狗腿度: {max_dogleg:.2f}°/30m | "
                                 f"平均狗腿度: {avg_dogleg:.2f}°/30m | "
                                 f"高狗腿度点: {high_dogleg_count}")
            control_layout.addWidget(stats_label)

        control_layout.addStretch()

        # 添加切换按钮
        self.toggle_button = QPushButton("切换到井斜深视图")
        self.toggle_button.clicked.connect(self.toggle_view_mode)
        control_layout.addWidget(self.toggle_button)

        layout.addLayout(control_layout)

        # 创建 3D 绘图部件
        self.plot_widget = ThreeDPlotWidget(self, width=12, height=8)
        layout.addWidget(self.plot_widget)

        # 绘制图形
        self.update_plot()

    def toggle_view_mode(self):
        """切换视图模式"""
        self.use_dogleg_severity = not self.use_dogleg_severity

        if self.use_dogleg_severity:
            self.toggle_button.setText("切换到井斜深视图")
        else:
            self.toggle_button.setText("切换到狗腿度视图")

        self.update_plot()

    def update_plot(self):
        """更新绘图"""
        self.plot_widget.plot_well_trajectory_with_dogleg(self.df, self.use_dogleg_severity)


def create_3d_well_plot(df):
    """创建并显示 3D 井眼轨迹图"""
    app = QApplication(sys.argv)
    window = MainWindow(df)
    window.show()
    sys.exit(app.exec_())


def calculate_displacement(df):
    """
    计算视位移、北移和东移
    返回:
    df: 添加了视位移、北移、东移三列的新DataFrame
    """

    # 初始化结果列
    df['视位移'] = 0.0
    df['北移'] = 0.0
    df['东移'] = 0.0
    df['狗腿角'] = 0.0  # 新增狗腿角列
    df['狗腿度'] = 0.0  # 新增狗腿度列

    # 计算视位移（从第二行开始）
    for i in range(1, len(df)):
        # 获取当前行和前一行数据
        A_prev = df.loc[i - 1, '井斜深']  # A4
        A_curr = df.loc[i, '井斜深']  # A5
        B_prev = df.loc[i - 1, '井斜角']  # B4
        B_curr = df.loc[i, '井斜角']  # B5
        E_prev = df.loc[i - 1, '视位移']  # E4

        # 角度转弧度
        angle_diff_rad = (B_curr - B_prev) * np.pi / 180
        avg_angle_rad = (B_prev + B_curr) / 2 * np.pi / 180

        # 计算视位移
        df.loc[i, '视位移'] = (E_prev +
                               (1 - (angle_diff_rad ** 2 / 24)) *
                               (A_curr - A_prev) *
                               np.sin(avg_angle_rad))

    # 计算北移、东移、狗腿角和狗腿度（从第二行开始）
    for i in range(1, len(df)):
        # 获取当前行和前一行数据
        A_prev = df.loc[i - 1, '井斜深']  # A3
        A_curr = df.loc[i, '井斜深']  # A4
        B_prev = df.loc[i - 1, '井斜角']  # B3
        B_curr = df.loc[i, '井斜角']  # B4
        C_prev = df.loc[i - 1, '方位角']  # C3
        C_curr = df.loc[i, '方位角']  # C4
        F_prev = df.loc[i - 1, '北移']  # F3
        G_prev = df.loc[i - 1, '东移']  # G3

        # 计算狗腿角 - 根据Excel公式转换
        # Excel公式: =ACOS(COS(B3/180*PI())*COS(B4/180*PI())+SIN(B3/180*PI())*SIN(B4/180*PI())*COS((C4-C3)/180*PI()))/PI()*180
        B_prev_rad = B_prev * np.pi / 180  # B3/180*PI()
        B_curr_rad = B_curr * np.pi / 180  # B4/180*PI()
        C_diff_rad = (C_curr - C_prev) * np.pi / 180  # (C4-C3)/180*PI()

        cos_dogleg = (np.cos(B_prev_rad) * np.cos(B_curr_rad) +
                      np.sin(B_prev_rad) * np.sin(B_curr_rad) * np.cos(C_diff_rad))

        # 处理浮点精度问题，确保cos值在[-1, 1]范围内
        cos_dogleg = np.clip(cos_dogleg, -1.0, 1.0)

        dogleg_angle_rad = np.arccos(cos_dogleg)  # ACOS(...)
        dogleg_angle_deg = dogleg_angle_rad * 180 / np.pi  # /PI()*180

        df.loc[i, '狗腿角'] = dogleg_angle_deg

        # 计算狗腿度 - 根据Excel公式转换
        # Excel公式: =30*H4/(A4-A3)
        # H4代表狗腿角（我们刚刚计算出来的），A4-A3代表井斜深差
        md_difference = A_curr - A_prev
        if md_difference > 0:
            dogleg_severity = 30 * df.loc[i, '狗腿角'] / md_difference
        else:
            dogleg_severity = 0

        df.loc[i, '狗腿度'] = dogleg_severity

        # 角度转弧度（用于位移计算）
        B_diff_rad = (B_curr - B_prev) * np.pi / 180
        C_diff_rad = (C_curr - C_prev) * np.pi / 180
        B_avg_rad = (B_prev + B_curr) / 2 * np.pi / 180
        C_avg_rad = (C_prev + C_curr) / 2 * np.pi / 180

        # 计算修正系数
        correction_factor = 1 - (B_diff_rad ** 2 + C_diff_rad ** 2) / 24

        # 计算北移
        df.loc[i, '北移'] = (F_prev +
                             correction_factor *
                             (A_curr - A_prev) *
                             np.sin(B_avg_rad) *
                             np.cos(C_avg_rad))

        # 计算东移
        df.loc[i, '东移'] = (G_prev +
                             correction_factor *
                             (A_curr - A_prev) *
                             np.sin(B_avg_rad) *
                             np.sin(C_avg_rad))
    print(df)
    return df


# 使用示例
if __name__ == "__main__":
    # 假设 df 是包含计算结果的 DataFrame
    result_df = pd.read_csv('/Users/cuitengfei/xishida/project/deep_well/data/demo.txt', sep='\t', encoding='utf-8')
    result_df = calculate_displacement(result_df)
    create_3d_well_plot(result_df)